import logging
import os

import streamlit as st
from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

# 配置日志记录
logging.basicConfig(level=logging.DEBUG)
'''
streamlit run class03/p3_3_4streamlit.py
'''

@st.cache_resource(show_spinner=False)
def load_data():
    """加载文档数据并创建索引，使用缓存提高性能"""
    # 初始化模型对象
    llm_dash = DashScope(
        model="qwen-max",  # 使用通义千问Max模型
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从环境变量获取API密钥
        temperature=0.3,  # 控制生成结果的随机性
        system_prompt="""作为手机选购助手，我将帮助用户选择合适手机。
              仅根据提供的文档内容进行回答。
              如果问题与文档无关，请回答：抱歉，这个话题我不擅长。"""
    )

    # 使用本地嵌入模型（节省成本）
    embed_model_zh = HuggingFaceEmbedding(
        model_name="BAAI/bge-small-zh-v1.5",  # 中文优化的嵌入模型
        device="cpu"  # 使用CPU运行
    )

    # 显示加载进度提示
    with st.spinner(text="正在加载文档数据，请稍候..."):
        data_dir = 'D:/pythonRoom/llm_llama/data2'  # 文档数据目录
        reader = SimpleDirectoryReader(input_dir=data_dir, recursive=True)
        docs = reader.load_data()  # 加载文档数据

        # 配置全局设置
        Settings.embed_model = embed_model_zh
        Settings.llm = llm_dash

        # 创建向量索引
        index = VectorStoreIndex.from_documents(docs)
        return index


# 主程序入口
if __name__ == "__main__":
    # 设置页面标题和图标
    st.set_page_config(
        page_title="手机选购助手",
        page_icon="📱",
        layout="centered"
    )

    # 初始化聊天引擎
    if "chat_engine" not in st.session_state:
        with st.spinner("正在初始化AI助手..."):
            # 加载数据并创建索引
            vector_index = load_data()
            # 创建查询引擎
            query_engine = vector_index.as_query_engine()
            # 创建聊天引擎（带上下文压缩功能）
            chat_engine = CondenseQuestionChatEngine.from_defaults(
                query_engine=query_engine,
                verbose=True  # 输出详细日志
            )
            st.session_state.chat_engine = chat_engine

    # 初始化聊天消息历史
    if "messages" not in st.session_state:
        st.session_state.messages = [
            {"role": "assistant", "content": "您好！我是手机选购助手，请问有什么可以帮您？"}
        ]

    # 显示所有历史聊天消息
    for message in st.session_state.messages:
        with st.chat_message(message["role"], avatar="📱" if message["role"] == "assistant" else "👤"):
            st.markdown(message["content"])

    # 获取用户输入
    if prompt := st.chat_input("请输入您的手机相关问题..."):
        # 添加用户消息到历史记录并显示
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user", avatar="👤"):
            st.markdown(prompt)

        # 生成助手回复
        with st.chat_message("assistant", avatar="📱"):
            with st.spinner("正在思考..."):
                # 获取聊天引擎的回复
                response = st.session_state.chat_engine.chat(prompt)
                # 显示回复内容
                st.markdown(response.response)

            # 添加助手回复到历史记录
            st.session_state.messages.append(
                {"role": "assistant", "content": response.response}
            )
